from data_type import *
from decoder import type, mask, sign_expend
from MMU import memory, MMU
from system import sys_call, ecall


# exec ================================================================
def _nop_(A: bit_64, B: bit_64, job: type):
    job.wb = True
    return bit_64()


def _add_(A: bit_64, B: bit_64, job: type):
    return bit_64(A.val + B.val)


def _sub_(A: bit_64, B: bit_64, job: type):
    return bit_64(A.val - B.val)


def _sra_(A: bit_64, B: bit_64, job: type):
    shamt = B.val & 0b111111
    if A.val >> 63:
        return (A.val >> shamt) | (mask(shamt) << (64 - shamt))
    else:
        return bit_64(A.val >> shamt)


def _srl_(A: bit_64, B: bit_64, job: type):
    return bit_64(A.val >> (B.val & 0b111111))


def _and_(A: bit_64, B: bit_64, job: type):
    return bit_64(A.val & B.val)


def _xor_(A: bit_64, B: bit_64, job: type):
    return bit_64(A.val ^ B.val)


def _or_(A: bit_64, B: bit_64, job: type):
    return bit_64(A.val | B.val)


def _sll_(A: bit_64, B: bit_64, job: type):
    return bit_64(A.val << (B.val & 0b111111))


def _mul_(A: bit_64, B: bit_64, job: type):
    return bit_64(A.val * B.val)


def _mulh_(A: bit_64, B: bit_64, job: type):
    if job.unsigned:
        return bit_64((A.val * B.val) >> 64)
    else:
        return bit_64((sign_expend(A.val, 64) * sign_expend(B.val, 64)) >> 64)


def _div_(A: bit_64, B: bit_64, job: type):
    if job.unsigned:
        return bit_64((A.val // B.val))
    else:
        return bit_64((sign_expend(A.val, 64) // sign_expend(B.val, 64)))


def _rem_(A: bit_64, B: bit_64, job: type):
    if job.unsigned:
        return bit_64(A.val % B.val)
    else:
        return bit_64(sign_expend(A.val, 64) % sign_expend(B.val, 64))


def _slt_(A: bit_64, B: bit_64, job: type):
    if job.unsigned:
        return bit_64(A.val < B.val)
    else:
        return bit_64(sign_expend(A.val, 64) < sign_expend(B.val, 64))


# ===================================================================================


lib_exec = {"add": _add_, "sub": _sub_, "sra": _sra_, "srl": _srl_, "and": _and_, "xor": _xor_, "or": _or_,
            "sll": _sll_, "": _nop_,
            "mul": _mul_, "mulh": _mulh_, "div": _div_, "rem": _rem_, "slt": _slt_}
cut = (0xffffffffffffffff, 0xffffffff, 0xffff, 0xff)
put = (64, 32, 16, 8)


def EX(job: type, rs1_r: bit_64, rs2_r: bit_64, imm: bit_64, rd, pc: bit_64, width, mem: memory, mmu: MMU, call: ecall):
    jump = False
    wb = bit_64()
    pc_next = bit_64()
    addr = bit_64()
    content = bit_64()
    if job.do == 0:
        ops = lib_exec[job.name]
        if job.imm:
            temp = ops(rs1_r, imm, job)
        else:
            temp = ops(rs1_r, rs2_r, job)
        wb = bit_64(sign_expend(temp.val & cut[job.size], put[job.size]))
        job.wb = not job.wb
#        print("EX: %d %s (%d/%d) = %d" % (rs1_r.val, job.name, rs2_r.val, imm.val, wb.val))
    elif job.do == 1 or job.do == 2:
        addr = bit_64(rs1_r.val + imm.val)
        if job.do == 1:
            content = bit_64(rs2_r.val)
        else:
            job.wb = True
#        print("EX: addr = %d + %d = %d, content = %d" % (rs1_r.val, imm.val, addr.val, content.val))
    elif job.do == 3:
        if job.name[0] == 'j':
            wb = bit_64(pc.val + (2 if width == 16 else 4))
            job.wb = True
            jump = True
            if job.name == "jal":
                pc_next = bit_64(imm.val + pc.val)
            else:
                pc_next = bit_64((rs1_r.val + imm.val) & 0xfffffffffffffffe)
        else:
            pc_next = bit_64(imm.val + pc.val)
            if job.name == "beq":
                if rs1_r.val == rs2_r.val:
                    jump = True
            elif job.name == "bne":
                if rs1_r.val != rs2_r.val:
                    jump = True
            elif job.name == "blt":
                if job.unsigned:
                    if rs1_r.val < rs2_r.val:
                        jump = True
                else:
                    if sign_expend(rs1_r.val, 64) < sign_expend(rs2_r.val, 64):
                        jump = True
            elif job.name == "bge":
                if job.unsigned:
                    if rs1_r.val >= rs2_r.val:
                        jump = True
                else:
                    if sign_expend(rs1_r.val, 64) >= sign_expend(rs2_r.val, 64):
                        jump = True
#        if jump:
#            print("EX: jump to 0x%x load %d to x%d" % (pc_next.val, wb.val, rd))
#            print("test test")
#        else:
#            print("EX: NOP")
    elif job.do == 4:
        if job.name == "lui":
            wb = bit_64(imm.val)
            job.wb = True
#            print("EX: x%d = %d" % (rd, wb.val))
        elif job.name == "auipc":
            wb = bit_64(imm.val + pc.val)
            job.wb = True
#            print("EX: x%d = %d" % (rd, wb.val))
        elif job.name == "ecall":
            job.wb = True
            rd = 10
            wb = bit_64(sys_call(call, mem, mmu))
#            print("EX: ecall -> %d" % call.a[7].val)

        elif job.name == "nop":
            pass
        elif job.name == "csrrw":
            pass
        elif job.name == "csrrs":
            pass
        elif job.name == "cssrrsi":
            pass

    return jump, pc_next, wb, addr, content, rd, job
